import React, {useState, useEffect, useMemo} from "react";
import {Tree, Checkbox, Button, Tooltip, Input} from "antd";
import type {TreeProps} from "antd";
import {SearchOutlined, ExpandOutlined, ShrinkOutlined} from "@ant-design/icons";

const {TreeNode} = Tree;
const {Search} = Input;

interface SlTreeProps extends Omit<TreeProps, "height"> {
    value?: React.Key[] | {
        checked: React.Key[];
        halfChecked: React.Key[]
    };
    showTool?: boolean;
    height?: string | number;
    onInput?: (value: any) => void;
    searchable?: boolean;
    expandAllByDefault?: boolean;
    showExpandAll?: boolean;
    showSelectAll?: boolean;
    showCheckStrictly?: boolean;
    fieldNames?: {
        title: string,
        key: string,
        children: string
    }
}

const SlTree: React.FC<SlTreeProps> = (
    {
        value,
        showTool = true,
        height = "400px",
        onInput,
        searchable = false,
        expandAllByDefault = true,
        showExpandAll = true,
        showSelectAll = true,
        showCheckStrictly = true,
        treeData,
        fieldNames = {
            title: "label",
            key: "id",
            children: "children"
        },
        ...treeProps
    }) => {
    // 状态管理
    const [expandedKeys, setExpandedKeys] = useState<React.Key[]>([]);
    const [autoExpandParent, setAutoExpandParent] = useState(true);
    const [searchValue, setSearchValue] = useState("");
    const [checkStrictly, setCheckStrictly] = useState(false);
    const [selectedKeys, setSelectedKeys] = useState<React.Key[]>([]);
    const [checkedKeys, setCheckedKeys] = useState<React.Key[] | {
        checked: React.Key[];
        halfChecked: React.Key[]
    }>([]);
    // 获取所有节点 key
    const getAllKeys = (data: any[]): React.Key[] => {
        return data.reduce((keys, node) => {
            keys.push(node[fieldNames.key]);
            if (node[fieldNames.children]) {
                keys.push(...getAllKeys(node[fieldNames.children]));
            }
            return keys;
        }, []);
    };

    // 所有节点 key
    const allKeys = useMemo(() => {
        return treeData ? getAllKeys(treeData) : [];
    }, [treeData]);

    useEffect(() => {
        if (value) {
            setCheckedKeys(value);
        }
    }, [value]);

    // 初始化展开所有节点
    useEffect(() => {
        if (expandAllByDefault && treeData) {
            setExpandedKeys(allKeys);
        }
    }, [treeData, expandAllByDefault, allKeys]);

    // 处理搜索
    const handleSearch = (searchValue: string) => {
        setSearchValue(searchValue);
        if (!searchValue) {
            setExpandedKeys([]);
            return;
        }

        const getParentKeys = (node: any, key: React.Key): React.Key[] => {
            if (!node || !node[fieldNames.children]) return [];
            return node[fieldNames.children].reduce((keys: React.Key[], child: any) => {
                if (child[fieldNames.title]?.toString().toLowerCase().includes(searchValue.toLowerCase())) {
                    keys.push(child[fieldNames.key]);
                    keys.push(...getParentKeys(child, child[fieldNames.key]));
                }
                return keys;
            }, []);
        };

        const expandedKeys = treeData?.reduce((keys: React.Key[], node: any) => {
            if (node[fieldNames.title]?.toString().toLowerCase().includes(searchValue.toLowerCase())) {
                keys.push(node[fieldNames.key]);
                keys.push(...getParentKeys(node, node[fieldNames.key]));
            }
            return keys;
        }, []) || [];

        setExpandedKeys(expandedKeys);
        setAutoExpandParent(true);
    };

    // 处理展开/折叠所有节点
    const toggleExpandAll = () => {
        if (expandedKeys.length > 0) {
            setExpandedKeys([]);
        } else {
            setExpandedKeys(allKeys);
        }
    };

    // 处理全选/取消全选
    const toggleSelectAll = () => {
        if (Array.isArray(checkedKeys) && checkedKeys.length === allKeys.length) {
            setCheckedKeys([]);
            onInput?.([]);
        } else {
            setCheckedKeys(allKeys);
            onInput?.(allKeys);
        }
    };

    // 处理父子联动切换
    const toggleCheckStrictly = () => {
        setCheckStrictly(!checkStrictly);
    };

    // 处理节点展开
    const handleExpand = (expandedKeys: React.Key[]) => {
        setExpandedKeys(expandedKeys);
        setAutoExpandParent(false);
    };

    // 处理节点选择
    const handleSelect = (selectedKeys: React.Key[], info: any) => {
        setSelectedKeys(selectedKeys);
        onInput?.(selectedKeys);
    };

    // 处理节点勾选
    const handleCheck = (keys: React.Key[] | { checked: React.Key[]; halfChecked: React.Key[] }, info: any) => {
        setCheckedKeys(keys);
        const keysToEmit = Array.isArray(keys)
            ? keys
            : keys?.checked;
        console.log(keys, '=========>', keysToEmit)
        onInput?.(keysToEmit);
    };

    // 渲染工具栏
    const renderTool = () => {
        return (
            <div className="sl-tree-toolbar" style={{marginBottom: 12}}>
                {showExpandAll && (
                    <Tooltip title={expandedKeys.length > 0 ? "折叠所有" : "展开所有"}>
                        <Button
                            type="primary"
                            icon={expandedKeys.length > 0 ? <ShrinkOutlined/> : <ExpandOutlined/>}
                            onClick={toggleExpandAll}
                        >{expandedKeys.length > 0 ? "折叠所有" : "展开所有"}</Button>
                    </Tooltip>
                )}

                {showSelectAll && (
                    <Tooltip title={
                        Array.isArray(checkedKeys) && checkedKeys.length === allKeys.length
                            ? "取消全选"
                            : "全选"
                    }>
                        <Button
                            type="primary"
                            className={"ml-1"}
                            onClick={toggleSelectAll}
                        >
                            {Array.isArray(checkedKeys) && checkedKeys.length === allKeys.length
                                ? "取消全选"
                                : "全选"}
                        </Button>
                    </Tooltip>
                )}

                {showCheckStrictly && (
                    <Checkbox
                        checked={checkStrictly}
                        onChange={toggleCheckStrictly}
                        style={{marginLeft: 8}}
                    >
                        父子联动
                    </Checkbox>
                )}
            </div>
        );
    };

    // 渲染搜索框
    const renderSearch = () => {
        if (!searchable) return null;

        return (
            <div style={{marginBottom: 12}}>
                <Search
                    placeholder="搜索节点"
                    allowClear
                    enterButton={<SearchOutlined/>}
                    onChange={(e) => handleSearch(e.target.value)}
                />
            </div>
        );
    };

    // 渲染树节点
    const renderTreeNodes = (data: any[]): React.ReactNode[] => {
        return data.map((item) => {
            const keyField = fieldNames.key ? fieldNames.key : "key";
            const titleField = fieldNames.title ? fieldNames.title : "title";
            const childrenField = fieldNames.children ? fieldNames.children : "children";
            const index = item[titleField]?.toString().toLowerCase().indexOf(searchValue.toLowerCase());
            const beforeStr = item[titleField]?.toString().substring(0, index);
            const matchStr = item[titleField]?.toString().substring(index, index + searchValue.length);
            const afterStr = item[titleField]?.toString().substring(index + searchValue.length);

            const title =
                index > -1 ? (
                    <span>
                        {beforeStr}
                        <span style={{color: "#f50"}}>{matchStr}</span>
                        {afterStr}
                    </span>
                ) : (
                    <span>{item[titleField]}</span>
                );

            return (
                <TreeNode
                    key={item[keyField]}
                    title={title}
                    data={item}
                    {...item}
                >
                    {item[childrenField] && renderTreeNodes(item[childrenField])}
                </TreeNode>
            );
        });
    };

    return (
        <div className="sl-tree-container">
            {showTool && renderTool()}
            {renderSearch()}

            <div style={{height, overflow: "auto"}}>
                <Tree
                    {...treeProps}
                    fieldNames={fieldNames}
                    expandedKeys={expandedKeys}
                    autoExpandParent={autoExpandParent}
                    selectedKeys={selectedKeys}
                    checkedKeys={checkedKeys}
                    checkStrictly={checkStrictly}
                    onExpand={handleExpand}
                    onSelect={handleSelect}
                    onCheck={handleCheck}
                >
                    {treeData ? renderTreeNodes(treeData) : null}
                </Tree>
            </div>
        </div>
    );
};

export default SlTree;